package com.autonavi.yunda.yunji.common.thread;

import org.jetbrains.annotations.NotNull;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.util.Assert;

import java.util.concurrent.*;

/**
 * @author cizhong.wcs
 * @date 2022/6/21 10:43
 */
public class TraceableThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {
    private int queueCapacity = Integer.MAX_VALUE;

    private TraceableThreadPoolExecutor threadPoolExecutor;

    public int getQueueCapacity() {
        return queueCapacity;
    }

    @Override
    public void setQueueCapacity(int queueCapacity) {
        this.queueCapacity = queueCapacity;
    }

    @NotNull
    @Override
    protected ExecutorService initializeExecutor(@NotNull ThreadFactory threadFactory,
                                                 @NotNull RejectedExecutionHandler rejectedExecutionHandler) {
        BlockingQueue<Runnable> queue = createQueue(this.queueCapacity);
        threadPoolExecutor = new TraceableThreadPoolExecutor(getCorePoolSize(), getMaxPoolSize(), getKeepAliveSeconds(),
                TimeUnit.SECONDS, queue, threadFactory, rejectedExecutionHandler);
        return threadPoolExecutor;
    }

    @NotNull
    @Override
    public ThreadPoolExecutor getThreadPoolExecutor() throws IllegalStateException {
        Assert.state(this.threadPoolExecutor != null, "ThreadPoolTaskExecutor not initialized");
        return this.threadPoolExecutor;
    }

    public synchronized int getTotalTasks() {
        return threadPoolExecutor.getTotalTasks();
    }

    public synchronized double getAverageTaskTime() {
        return threadPoolExecutor.getAverageTaskTime();
    }

}
